/***************************************
@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@!!!@@@@@@@@@@@@!!!!@@@@@@@@@@@%!!!@@@@@@@@@@@!!!!!!!!!!@@
@@!!!!!@@@@@@@@@%%!!!!!@@@@@@@@@%%!!!!@@@@@@@@@!!!!!!!!!!@@
@@@!!!!!@@@@@@@%%%%!!!!!@@@@@@@%%%!!!!!@@@@@@@@@@@@@!!!!!@@
@@@@!!!!!@@@@@@%%%@@!!!!!@@@@@%%%%@!!!!!@@@@@@@@@@@@!!!!!@@
@@@@@!!!!!@@@@%%%@@@@!!!!@@@@%%%%@@@!!!!!@@@@@@@@@@@!!!!!@@
@@@@@@!!!!!@@%%%%@@@@!!!!!@@@%%%@@@@@!!!!!@@@@@@@@@@!!!!!@@
@@@@@@!!!!!@%%%%@@@@@@!!!!!@%%%$@@@@@@!!!!@@@@@@@@@@!!!!!@@
@@@@@@@!!!!!%%%@@@@@@@@!!!!!%%%@@@@@@@@!!!!@@@@@@@@!!!!!!@@
@@@@@@@@!!!!!%@@@@@@@@@@!!!!!%@@@@@@@@@!!!!!!!!!!!!!!!!!!@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@

WMJ战队深度学习目标检测TensorRT推理后端
依赖库：                 TensorRT-8.6，CUDA-12.1，NPP库
目前支持使用的模型：       yolox改四点模型，大符/装甲板识别
更新时间：               2024.02.01
***************************************/

#pragma once
#include "InferenceBackend.hpp"
#include <dirent.h>
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "logging.h"
#include <npp.h>

using namespace nvinfer1;

#define CHECK(status) \
    do\
    {\
        auto ret = (status);\
        if (ret != 0)\
        {\
            std::cerr << "Cuda failure: " << ret << std::endl;\
            abort();\
        }\
    } while (0)



namespace wmj
{
    class TRTBackend : public InferenceBackendBase
    {
    public:
        TRTBackend(){};
        TRTBackend(const DetectorParams &params);
        ~TRTBackend();
        bool doInference(const cv::Mat &inputImg);
        const float* getOutput();
        enum BackendType type_ = TRT;
        float* m_prob;
    private:
        void doInference(IExecutionContext& context, uchar* input, float* output, void* buffers[], void* gpu_img,NppStreamContext &nppCtx,cudaStream_t &stream);
        Logger gLogger;
        IRuntime* m_runtime = nullptr;
        ICudaEngine* m_engine = nullptr;
        IExecutionContext* m_context = nullptr;
        cudaStream_t m_stream;
        NppStreamContext nppCtx;
        void* m_buffers[2];
        void  *m_gpu_img;
        int DEVICE = 0;

        int inputIndex,outputIndex;
        int input_size,output_size,gpu_img_size;

        bool m_output_flag;   
    };


    TRTBackend::TRTBackend(const DetectorParams &params)
    {
        type_ = TRT;
        assert(params.modelFormat == "trt");
        std::string model_name = MODEL_DIR + params.modelName + ".engine";
        cudaSetDevice(0);
        char *trtModelStream{nullptr};
        size_t size{0};
        //*********** read model *************//
        std::ifstream file(model_name, std::ios::binary);
        if (file.good()) 
        {
            file.seekg(0, file.end);
            size = file.tellg();
            file.seekg(0, file.beg);
            trtModelStream = new char[size];
            assert(trtModelStream);
            file.read(trtModelStream, size);
            file.close();
        }
        //*********** read model *************//

        m_runtime = createInferRuntime(gLogger);
        assert(m_runtime != nullptr);
        m_engine = m_runtime->deserializeCudaEngine(trtModelStream, size);
        assert(m_engine != nullptr); 
        m_context = m_engine->createExecutionContext();

        assert(m_context != nullptr);
        delete[] trtModelStream;

        auto in_dims = m_engine -> getTensorShape("input");
        auto out_dims = m_engine -> getTensorShape("output");
        output_size = 1;
        for(int j=0;j<out_dims.nbDims;j++) {
            output_size *= out_dims.d[j];
        }
        m_prob  = new float[output_size];
        input_width = in_dims.d[3];
        input_height = in_dims.d[2];
        objectLength = out_dims.d[2];
        maxProposalCount = out_dims.d[1];
        assert(in_dims.d[0] == out_dims.d[0]);
        batch_size = in_dims.d[0];
        int input_channel = in_dims.d[1];

        assert(m_engine->getNbBindings() == 2);
        inputIndex = m_engine->getBindingIndex("input");
        assert(m_engine->getBindingDataType(inputIndex) == nvinfer1::DataType::kFLOAT);
        assert(m_engine -> getTensorFormat("input") == nvinfer1::TensorFormat::kHWC);

        outputIndex = m_engine->getBindingIndex("output");
        assert(m_engine->getBindingDataType(outputIndex) == nvinfer1::DataType::kFLOAT);
        input_size = batch_size * input_channel * input_width * input_height * sizeof(float);
        gpu_img_size = input_channel * input_width * input_height * sizeof(uchar);
        output_size *= sizeof(float);
        CHECK(cudaMalloc(&m_buffers[inputIndex], input_size));
        CHECK(cudaMalloc(&m_buffers[outputIndex], output_size));

        CHECK(cudaMalloc(&m_gpu_img, gpu_img_size));

        CHECK(cudaStreamCreate(&m_stream));
        nppCtx.hStream = m_stream;
    }

    bool TRTBackend::doInference(const cv::Mat &inputImg)
    {
        if(inputImg.cols != input_width || inputImg.rows != input_height)
        {
            std::cout << "input Img size not match!! Model input size: " 
            << input_width << "x" << input_height << " But input Image size is: "
            << inputImg.cols << "x" << inputImg.rows << std::endl;
            return false;
        }
        doInference(*m_context, inputImg.data, m_prob, m_buffers, m_gpu_img, nppCtx, m_stream);
        return true;
    }

    const float* TRTBackend::getOutput()
    {
        if(m_output_flag)
        {
            m_output_flag = false;
            return m_prob;
        }
        else
        {
            std::cout << "output_flag set FALSE, inference result maybe out of time, return NULL" << std::endl;
            return nullptr;
        }

    }

    void TRTBackend::doInference(IExecutionContext& context, uchar* input, float* output, void* buffers[], void* gpu_img,NppStreamContext &nppCtx,cudaStream_t &stream)
    {
        CHECK(cudaMemcpyAsync(gpu_img, input, gpu_img_size, cudaMemcpyHostToDevice,stream));
        nppiConvert_8u32f_C3R_Ctx((Npp8u*)gpu_img, input_width * 3, (Npp32f*)buffers[inputIndex], input_width * 3 * sizeof(float), NppiSize{input_width, input_height},nppCtx);
        context.enqueueV2(buffers,stream,nullptr);
        CHECK(cudaMemcpyAsync(output, buffers[outputIndex], output_size, cudaMemcpyDeviceToHost, stream));
        cudaStreamSynchronize(stream);
        m_output_flag = true;
    }

    TRTBackend::~TRTBackend()
    {
        delete[] m_prob;

        cudaStreamDestroy(m_stream);
        cudaFree(m_buffers[outputIndex]);
        cudaFree(m_buffers[inputIndex]);
        cudaFree(m_gpu_img);

        delete m_context;
        delete m_engine;
    }

}

